from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.chains.sequential import SimpleSequentialChain
from langchain_community.llms.baidu_qianfan_endpoint import QianfanLLMEndpoint

template_name = """
你是一个非常棒的起名字专家，你非常擅长给新生儿们起好听的名字，用中文给出3种回答。

这里是问题：
{input}
"""

prompt_template_name = PromptTemplate(input_variables=["input"], template=template_name)
llm = QianfanLLMEndpoint(
    qianfan_ak='ae7eoePCWBS8nBTQv7vbmyxm',
    qianfan_sk='tvxwV0Emkp9sxuU3neDfF77Drhl7TWUa'
)
chain_name = LLMChain(llm = llm, prompt=prompt_template_name)

template_judge = """
你是一个名字评论专家，你非常擅长给名字做分析，并且判断哪个名字更好。

名字分析：
{name}
对上面的名字进行分析，判断哪个更好。
"""
prompt_template_judge = PromptTemplate.from_template(template_judge)
chain_judge = LLMChain(llm = llm, prompt=prompt_template_judge)

chain_sequence = SimpleSequentialChain(chains=[chain_name, chain_judge], verbose=True)
review = chain_sequence.invoke("给我一个很土但是听起来好养活的男孩小名")